import random
import time
from typing import List
from collections import deque
import matplotlib.pyplot as plt
import matplotlib as mpl

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        nums.sort(reverse=True)
        return nums[k-1]
    
    def findKthLargestByStack(self, nums: List[int], k: int) -> int:
        stack = deque()
        for num in nums:
            if len(stack) < k:
                stack.append(num)
                stack = deque(sorted(stack, reverse=True))
            elif num > stack[-1]:
                stack.pop()
                stack.append(num)
                stack = deque(sorted(stack, reverse=True))
        return stack[-1]

def run_comparison():
    # 设置不同的数组规模
    sizes = [1000, 5000, 10000, 20000, 25000, 50000, 100000, 150000]
    sort_times = []
    stack_times = []
    k_ratio = 0.1  # k值设为数组大小的10%

    sol = Solution()

    # 对每个规模进行测试
    for n in sizes:
        k = max(1, int(n * k_ratio))
        # 生成随机数组
        nums = [random.randint(0, 100000) for _ in range(n)]

        # 测试排序方法
        start_time = time.time()
        sol.findKthLargest(nums.copy(), k)
        sort_times.append((time.time() - start_time) * 1000)

        # 测试压缩栈方法
        start_time = time.time()
        sol.findKthLargestByStack(nums.copy(), k)
        stack_times.append((time.time() - start_time) * 1000)

        print(f"完成规模 {n} 的测试")

    # 创建图表并设置大小
    plt.figure(figsize=(12, 7))

    # 绘制曲线
    plt.plot(sizes, sort_times, 'b-o', label='排序方法', linewidth=2)
    plt.plot(sizes, stack_times, 'r-s', label='压缩栈方法', linewidth=2)

    # 设置图表属性（使用中文标签）
    plt.title('查找第K大元素的两种方法效率对比', fontsize=14)
    plt.xlabel('数组规模 (n)', fontsize=12)
    plt.ylabel('运行时间 (毫秒)', fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend(fontsize=10)

    # 定义需要标注的关键点（在0-20000范围内只标注1000和10000）
    key_points = {1000, 10000}

    # 分别对两条线进行标注，避免贴错
    for i, size in enumerate(sizes):
        if size <= 20000 and size not in key_points:
            continue  # 跳过非关键点

        # 标注排序方法（蓝色线）
        sort_time = sort_times[i]
        plt.annotate(f'{sort_time:.2f}ms', 
                    (size, sort_time), 
                    textcoords="offset points", 
                    xytext=(0, 10), 
                    ha='center',
                    fontsize=9,
                    color='blue')

        # 标注压缩栈方法（红色线）
        stack_time = stack_times[i]
        plt.annotate(f'{stack_time:.2f}ms', 
                    (size, stack_time), 
                    textcoords="offset points", 
                    xytext=(0, -15), 
                    ha='center',
                    fontsize=9,
                    color='red')

    # 调整布局
    plt.tight_layout()

    # 保存图表（设置dpi以提高清晰度）
    plt.savefig('效率比较.png', dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    run_comparison()